#!/usr/bin/env python
# coding=utf-8
# @Software: PyCharm
#author GWC
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from scipy.signal import butter, filtfilt
import mne
# from plugin import Plugin

from abc import ABC, abstractmethod


class PluginInterface(ABC):
    @abstractmethod
    def execute(self,data):
        pass


#电极定位
class Position62Plugin(PluginInterface):
    def execute(self,data):
        print("执行插件 电极定位62导")
        # 将获取的电极位置信息修改并补充缺失的电极位置，整合为1020.xlsx
        data1020 = pd.read_excel('1020.xlsx', index_col=0)
        channels1020 = np.array(data1020.index)
        value1020 = np.array(data1020)
        # 将电极通道名称和对应三维坐标位置存储为字典形式
        list_dic = dict(zip(channels1020, value1020))
        # print(list_dic)
        # 封装为MNE的格式，参考原biosemi的存储格式
        montage_1020 = mne.channels.make_dig_montage(ch_pos=list_dic,
                                                     nasion=[5.27205792e-18, 8.60992398e-02, -4.01487349e-02],
                                                     lpa=[-0.08609924, -0., -0.04014873],
                                                     rpa=[0.08609924, 0., -0.04014873])
        """
        #电极定位，绘制电极位置
        """
        # int1020_montage = mne.channels.make_standard_montage('standard_1020')
        position_data=data.set_montage(montage_1020, on_missing='warn')
        return position_data

class Position64Plugin(PluginInterface):
    def execute(self,data):
        print("执行插件 电极定位64导")
        int1020_montage = mne.channels.make_standard_montage('standard_1020')
        position_data=data.set_montage(int1020_montage, on_missing='warn')
        return position_data



#插值坏导
class PassPlugin(PluginInterface):

    def execute(self, data):
        print("执行插件 插值坏导")
        scalings = {'eeg': 50}
        raw_cropped = data.copy()
        raw_cropped.plot(scalings=scalings, block=True,title='请检查并选中坏导')  # 定义坏导
        plt.show()
        badflag = False
        if raw_cropped.info['bads']:
            print('已选择坏导: ', raw_cropped.info['bads'], '开始进行插值')
            badflag = True
        else:
            print('无坏导，跳过插值')
            return raw_cropped
        if badflag:
            # raw_pass.load_data()
            raw_pass =raw_cropped.interpolate_bads(exclude=[])
            raw_pass.plot(scalings=scalings, block=True,title='坏导插值完成，如无误请关闭窗口')
            plt.show()

        return raw_pass


#重参考
class RereferencePlugin(PluginInterface):

    def execute(self, data):
        scalings = {'eeg': 50}
        raw_ref=data.set_eeg_reference(ref_channels='average')
        raw_ref.plot(scalings=scalings, block=True,title='重参考完成，无误请关闭窗口')
        return raw_ref

#过滤插件
class FilterPlugin(PluginInterface):
    def __init__(self, l_freq, h_freq, method):
        self.l_freq = l_freq
        self.h_freq = h_freq
        self.method=method

    def execute(self,data):
        print("执行插件 滤波降噪")
        scalings = {'eeg': 50}
        raw_filter = data.copy()
        raw_filter = raw_filter.filter(self.l_freq, self.h_freq, method=self.method)
        raw_filter = raw_filter.notch_filter(self.h_freq)
        # a = raw_filter.plot_psd(self.l_freq, self.h_freq, sphere=(0, 0, 0, 0.11))
        # a.savefig("images/" + "PSD.png")
        # plt.show(block=True)
        raw_filter.plot(scalings=scalings,block=True, title='滤波完成，准备ICA，无误请关闭窗口')

        return raw_filter


#去除伪迹
class ICAPlugin(PluginInterface):

    def execute(self, data):
        # ICA去除伪迹
        scalings = {'eeg': 50}
        ica = mne.preprocessing.ICA(n_components=30,
                                    noise_cov=None,
                                    random_state=None,
                                    method='infomax',
                                    fit_params=None,
                                    max_iter='auto',
                                    allow_ref_meg=False,
                                    verbose=True)
        raw_recons = data.copy()
        ica.fit(raw_recons)
        # raw_recons.load_data()
        ica.plot_components()
        ica.plot_sources(raw_recons,block=True, show_scrollbars=False, title='请选择需要去除的成分')
        plt.show()
        raw_recons = ica.apply(raw_recons)
        data.plot(scalings=scalings, title='ICA处理前, 确认请关闭')
        raw_recons.plot(scalings=scalings,block=True, title='ICA处理后, 确认请关闭')
        plt.show(block=True)
        # print(ica)
        return raw_recons







class PluginManager:
    def __init__(self):
        self.plugins = []

    def add_plugin(self, plugin):
        self.plugins.append(plugin)

    def execute_all(self,data):

        for plugin in self.plugins:
            data=plugin.execute(data)
        return data


